use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicI64, Ordering};

pub struct DebugAlloc {
    count: AtomicI64, // 每个实例拥有独立的计数器
}

impl DebugAlloc {
    pub const fn new() -> Self {
        Self {
            count: AtomicI64::new(0),
        }
    }

    pub fn init_count(&self) {
        self.count.store(0, Ordering::Relaxed);
    }

    pub fn get_count(&self) -> i64 {
        self.count.load(Ordering::Relaxed)
    }
}

unsafe impl GlobalAlloc for DebugAlloc {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        unsafe {
            let size = layout.size();
            let mem = System.alloc(layout);
            if !mem.is_null() {
                self.count.fetch_add(size as i64, Ordering::Relaxed);
            }
            mem
        }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unsafe {
            self.count
                .fetch_sub(layout.size() as i64, Ordering::Relaxed);
            System.dealloc(ptr, layout);
        }
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        unsafe {
            let old_size = layout.size();
            let new_ptr = System.realloc(ptr, layout, new_size);
            if !new_ptr.is_null() {
                self.count.fetch_sub(old_size as i64, Ordering::Relaxed);
                self.count.fetch_add(new_size as i64, Ordering::Relaxed);
            }
            new_ptr
        }
    }
}
